{-# LANGUAGE DeriveDataTypeable,BangPatterns,ScopedTypeVariables #-}
module Management where

import Types
import Random
import Constants

import Data.Array.Vector
import Data.Array.Vector.UArr
import Data.Array.ST
import Data.List (sortBy,foldl',lookup)
import Data.Word
import qualified Data.IntMap as IntMap
import Data.Maybe
import Data.Binary
import Data.Generics


import Control.Monad
import Control.Monad.ST
import Control.DeepSeq

import Statistics.RandomVariate as RV

import Debug.Trace

-- **** Types ****


-- *** BANK MANAGEMENT AND MANAGEMENT POLICIES ***

data Management = 
    SimpleM !Policy !Policy !PNum ![Policy] -- separate policy for deposits and loans - pnum is divdend size
    deriving (Ord,Eq,Data,Typeable)

mRegs (SimpleM _ _ _ p) = p

instance NFData Management where
   rnf _ = () -- should already be strict ...

-- XXX - BROKEN!
instance Binary Management where
  put (SimpleM a b c d) = put a >> put b >> put c >> put d
  get = get >>= \a -> get >>= \b -> get >>= \c -> get >>= \d -> return (SimpleM a b c d)


data Mutator 
    = Mut { polmut::ProbabilityDistribution, newnummut::ProbabilityDistribution }
    deriving (Show,Data,Typeable)

instance NFData Mutator where
  rnf (Mut a b) = rnf a `seq` rnf b

instance Binary Mutator where
  put (Mut a b) = put a >> put b
  get = get >>= \a -> get >>= \b -> return (Mut a b)

data PDirection = PLT | PGT
                deriving (Show,Ord,Eq,Data,Typeable)

instance Binary PDirection where
  put PLT = putWord8 0
  put PGT = putWord8 1
  get = do
    tag_ <- getWord8
    case tag_ of
      0 -> return PLT
      1 -> return PGT
      _ -> fail "no parse"


data BSSide = Asset | Liability
            deriving (Show,Ord,Eq,Data,Typeable)

instance Binary BSSide where
  put Asset = putWord8 0
  put Liability = putWord8 1
  get = do
    tag_ <- getWord8
    case tag_ of
      0 -> return Asset
      1 -> return Liability
      _ -> fail "no parse"


data PNum 
    = PCapitalRat |
      PLiqRat |
      PRate |
      PDefault | 
      PPrincipal |
      PMaturity |
      PBankCash |
      PCashRate |
      PSum !PNum !PNum |
      PDiff !PNum !PNum |
      PProd !PNum !PNum |
      PDiv !PNum !PNum |
      PBSSum !BSSide !Policy | -- Sum of a class of entries on the BS
      PNum !Double | -- for constants
      PNullNum | -- a null to be replaced with anything
      PNullVal -- a null to be replaced with a (PNum x)
    deriving (Show,Ord,Eq,Data,Typeable)


instance Binary PNum where
  put PCapitalRat = putWord8 0
  put PLiqRat = putWord8 1
  put PRate = putWord8 2
  put PDefault = putWord8 3
  put PPrincipal = putWord8 4
  put PMaturity = putWord8 5
  put PBankCash = putWord8 6
  put (PSum a b) = putWord8 7 >> put a >> put b
  put (PDiff a b) = putWord8 8 >> put a >> put b
  put (PBSSum a b) = putWord8 9 >> put a >> put b
  put (PNum a) = putWord8 10 >> put a
  put PNullNum = putWord8 11
  put (PProd a b) = putWord8 12 >> put a >> put b
  put (PDiv a b) = putWord8 13 >> put a >> put b
  put PCashRate = putWord8 14
  get = do
    tag_ <- getWord8
    case tag_ of
      0 -> return PCapitalRat
      1 -> return PLiqRat
      2 -> return PRate
      3 -> return PDefault
      4 -> return PPrincipal
      5 -> return PMaturity
      6 -> return PBankCash
      7 -> get >>= \a -> get >>= \b -> return (PSum a b)
      8 -> get >>= \a -> get >>= \b -> return (PDiff a b)
      9 -> get >>= \a -> get >>= \b -> return (PBSSum a b)
      10 -> get >>= \a -> return (PNum a)
      11 -> return PNullNum
      12 -> get >>= \a -> get >>= \b -> return (PProd a b)
      13 -> get >>= \a -> get >>= \b -> return (PDiv a b)
      14 -> return PCashRate
      _ -> fail "no parse"


data Policy 
    = PNullPolicy |
      PAccept |
      PReject |
      PIf !Policy !Policy !Policy | -- cond true false
      PCall !Int | -- call the policy in slot num  -- note the num is an integer and needs hard-coded mutate code
      PNumTest !PNum !PDirection !PNum  -- test the amount of capital
    deriving (Ord,Eq,Data,Typeable)

instance Binary Policy where
  put PNullPolicy = putWord8 0
  put PAccept = putWord8 1
  put PReject = putWord8 2
  put (PIf a b c) = putWord8 3 >> put a >> put b >> put c
  put (PCall a) = putWord8 4 >> put a
  put (PNumTest a b c) = putWord8 5 >> put a >> put b >> put c
  get = do
    tag_ <- getWord8
    case tag_ of
      0 -> return PNullPolicy
      1 -> return PAccept
      2 -> return PReject
      3 -> get >>= \a -> get >>= \b -> get >>= \c -> return (PIf a b c)
      4 -> get >>= \a -> return (PCall a)
      5 -> get >>= \a -> get >>= \b -> get >>= \c -> return (PNumTest a b c)
      _ -> fail "no parse"


instance (Show Policy) where
  show p = showp 4 0 p

instance (Show Management) where
  show (SimpleM p1 p2 d regs ) 
       = "Management;\n\n Dividend policy:\n " ++ (show d) ++ "\nDeposit Policy: \n" 
         ++ (show p1) ++ ("Lending policy: \n") ++ (show p2) ++ (" ** Registers ** \n") 
                ++ concat (Prelude.map (uncurry regShow) (zip [1..(length regs)] regs))
    where regShow i r = "\nReg " ++ (show i) ++" :\n" ++ (show r)

showp i t PAccept = (showtab i t) ++ "PAccept\n"
showp i t PReject = (showtab i t) ++ "PReject\n"
showp i t (PIf p1 p2 p3)
    = (showtab i t) ++ " PIf  \n" ++ (showp i (t + 6) p1) ++
      (showtab i t) ++ " Then \n" ++ (showp i (t + 6) p2) ++
      (showtab i t) ++ " Else \n" ++ (showp i (t + 6) p3)
showp i t (PNumTest n1 s n2) = (showtab i t) ++ " Test " ++ (show n1) ++ " " ++ (show s) ++ " " ++ (show n2) ++ "\n"
showp i t (PCall n) = (showtab i t) ++ "PCall " ++ (show n) ++ "\n" 
showp i t PNullPolicy = (showtab i t) ++ "PNullPolicy\n"

showtab i t = (foldl1 (++) (take (i + t) (repeat " "))) 


-- **** MUTATIONS ****
{-

Very simple genetic learner for now

2) Crossover

- find two nodes
- substitute the two subtrees
-- Later - for now we can get simple crossover going 
-- at the management level

-}

cp0 0 (SimpleM d1 d2 d3 m1) (SimpleM _ _ _ m2) = (SimpleM d1 d2 d3 (me1 m1 m2)) 
cp0 1 (SimpleM d1 d2 _ m1) (SimpleM _ _ d3 m2) = (SimpleM d1 d2 d3 (me2 m1 m2)) 
cp0 2 (SimpleM d1 _ d3 m1) (SimpleM _ d2 _ m2) = (SimpleM d1 d2 d3 (me1 m1 m2)) 
cp0 3 (SimpleM d1 _ _ m1) (SimpleM _ d2 d3 m2) = (SimpleM d1 d2 d3 (me2 m1 m2)) 
cp0 4 (SimpleM _ d2 d3 m1) (SimpleM d1 _ _ m2) = (SimpleM d1 d2 d3 (me1 m1 m2)) 
cp0 5 (SimpleM _ d2 _ m1) (SimpleM d1 _ d3 m2) = (SimpleM d1 d2 d3 (me2 m1 m2)) 
cp0 6 (SimpleM _ _ d3 m1) (SimpleM d1 d2 _ m2) = (SimpleM d1 d2 d3 (me1 m1 m2)) 
cp0 7 (SimpleM _ _ _ m1) (SimpleM d1 d2 d3 m2) = (SimpleM d1 d2 d3 (me2 m1 m2)) 
cp0 _ _ _ = (trace "cp0 error" undefined)

me1 [] _ = []
me1 _ [] = []
me1 (m1:m1s) (m2:m2s) = m1:(me2 m1s m2s)

me2 [] _ = []
me2 _ [] = []
me2 (m1:m1s) (m2:m2s) = m2:(me1 m1s m2s)

crossPoss = 7

crossOver1 :: Gen a -> Management -> Management -> ST a Management
crossOver1 g m1 m2 =
    do num <- randRange g 0 crossPoss
       return $ cp0 (num) m1 m2

crossOver :: Gen a -> [Management] -> ST a [Management]
crossOver g oldm =
    do (m1,m2) <- splitr g oldm
       mapM (uncurry (crossOver1 g)) (zip m1 m2)

splitr :: Gen a -> [Management] -> ST a ([Management],[Management])
splitr g ls =
    do rands <- rnums g (length ls)
       let blist = zip rands ls
       let slist = sortBy (\x y -> fst x `compare` fst y) blist
       let len = length slist
       let flist = map snd (if (len `mod` 2 == 0)
                            then slist
                            else tail slist)
       let halfway = (len - (len `mod` 2)) `div` 2
       return (splitAt halfway flist)

rnums :: Gen a -> Int -> ST a [Double]
rnums g n = replicateM n (RV.uniform g) -- (0,1]

data PNode 
     = PNP Policy
     | PNN PNum
     | PND PDirection
     | PNS BSSide
    deriving Show

-- Surely this is a job for a macro or Template Haskell!?!

maxn :: PNode -> Int
-- maxn x | (trace (show x) False) = undefined
maxn (PNP (PIf _ _ _)) = 3
maxn (PNP (PNumTest _ _ _)) = 3
maxn (PNP PAccept) = 0
maxn (PNP PReject) = 0
maxn (PNP (PCall _)) = 0
maxn (PNN (PSum _ _)) = 2
maxn (PNN (PDiff _ _)) = 2
maxn (PNN (PProd _ _)) = 2
maxn (PNN (PDiv _ _)) = 2
maxn (PNN (PBSSum _ _)) = 2
maxn (PNN _) = 0
maxn (PND _) = 0
maxn (PNS _) = 0

pick :: Int -> PNode -> PNode
-- pick i p | trace ("PICKING I: " ++ (show i) ++ " P: " ++ (show p)) False = undefined
pick 1 (PNP (PIf c _ _ ) ) = PNP c
pick 2 (PNP (PIf _ c _ ) ) = PNP c
pick 3 (PNP (PIf _ _ c ) ) = PNP c
pick _ (PNP PAccept) = (trace "pick PAccept" undefined)
pick _ (PNP PReject) = (trace "pick PReject" undefined)
pick _ (PNP (PCall n)) = (trace "pick PCall" undefined)
pick 1 (PNP (PNumTest d1 _ _)) = PNN d1
pick 2 (PNP (PNumTest _ d _)) = PND d
pick 3 (PNP (PNumTest _ _ d2)) = PNN d2
pick 1 (PNN (PSum n1 n2)) = PNN n1
pick 2 (PNN (PSum n1 n2)) = PNN n2
pick 1 (PNN (PDiff n1 n2)) = PNN n1 
pick 2 (PNN (PDiff n1 n2)) = PNN n2
pick 1 (PNN (PProd n1 n2)) = PNN n1 
pick 2 (PNN (PProd n1 n2)) = PNN n2
pick 1 (PNN (PDiv n1 n2)) = PNN n1 
pick 2 (PNN (PDiv n1 n2)) = PNN n2
pick 1 (PNN (PBSSum side pol)) = PNS side
pick 2 (PNN (PBSSum side pol)) = PNP pol
pick _ _ = (trace "pick other" undefined)

-- 
subst :: Int -> PNode -> PNode -> PNode
subst 1 (PNP new) (PNP (PIf _ c2 c3 ) ) = (PNP (PIf new c2 c3))
subst 2 (PNP new) (PNP (PIf c1 _ c3 ) ) = (PNP (PIf c1 new c3))
subst 3 (PNP new) (PNP (PIf c1 c2 _ ) ) = (PNP (PIf c1 c2 new))
subst 1 (PNN new) (PNP (PNumTest d1 d d2)) = (PNP (PNumTest new d d2))
subst 2 (PND new) (PNP (PNumTest d1 d d2)) = (PNP (PNumTest d1 new d2))
subst 3 (PNN new) (PNP (PNumTest d1 d d2)) = (PNP (PNumTest d1 d new))
subst 1 (PNN new) (PNN (PSum n1 n2)) = PNN (PSum new n2)
subst 2 (PNN new) (PNN (PSum n1 n2)) = PNN (PSum n1 new)
subst 1 (PNN new) (PNN (PDiff n1 n2)) = PNN (PDiff new n2)
subst 2 (PNN new) (PNN (PDiff n1 n2)) = PNN (PDiff n1 new)
subst 1 (PNN new) (PNN (PProd n1 n2)) = PNN (PProd new n2)
subst 2 (PNN new) (PNN (PProd n1 n2)) = PNN (PProd n1 new)
subst 1 (PNN new) (PNN (PDiv n1 n2)) = PNN (PDiv new n2)
subst 2 (PNN new) (PNN (PDiv n1 n2)) = PNN (PDiv n1 new)
subst 1 (PNS new) (PNN (PBSSum side pol)) = PNN (PBSSum new pol)
subst 2 (PNP new) (PNN (PBSSum side pol)) = PNN (PBSSum side new)

maxa :: PNode -> Int
maxa (PNP _) = 15
maxa (PNN _) = 9 -- NOTE this excludes BSSum and arithmetic for now which seems too hard for our little banks ...
maxa (PND _) = 2
maxa (PNS _) = 2

{-
-- PNullNum and PNullPolicy are placeholders to be replaced
--
-- needs to be changed to "alternate WEIGHT nodetype" where weight describes a prob distribution. 
-}

alternate :: Int -> PNode -> PNode
alternate 1 (PNP _) = (PNP PAccept)
alternate 2 (PNP _) = (PNP PReject)
alternate 3 (PNP _) = (PNP (PIf PNullPolicy PNullPolicy PNullPolicy)) 
alternate 4 (PNP _) = (PNP (PNumTest PNullNum PGT PNullNum)) 
alternate 5 (PNP _) = (PNP (PNumTest PNullNum PLT PNullNum)) 
alternate 6 (PNP _) = (PNP (PCall 1))
alternate 7 (PNP _) = (PNP (PCall 2))
alternate 8 (PNP _) = (PNP (PCall 3))
alternate 9 (PNP _) = (PNP (PCall 4))
alternate 10 (PNP _) = (PNP (PCall 5))
alternate 11 (PNP _) = (PNP (PCall 6))
alternate 12 (PNP _) = (PNP (PCall 7))
alternate 13 (PNP _) = (PNP (PCall 8))
alternate 14 (PNP _) = (PNP (PCall 9))
alternate 15 (PNP _) = (PNP (PCall 10))

alternate 1 (PNN _) = PNN PRate
alternate 2 (PNN _) = PNN PPrincipal
alternate 3 (PNN _) = PNN PDefault 
alternate 4 (PNN _) = PNN PMaturity 
alternate 5 (PNN _) = PNN PBankCash 
alternate 6 (PNN _) = PNN PCapitalRat 
alternate 7 (PNN _) = PNN PLiqRat
alternate 8 (PNN _) = PNN PCashRate
alternate 9 (PNN _) = PNN (PNullVal) 
--
alternate 10 (PNN _) = PNN (PProd PNullNum PNullNum) 
alternate 11 (PNN _) = PNN (PSum PNullNum PNullNum) 
alternate 12 (PNN _) = PNN (PDiff PNullNum PNullNum) 
alternate 13 (PNN _) = PNN (PDiv PNullNum PNullNum) 
alternate 14 (PNN _) = PNN (PBSSum Asset PNullPolicy) -- XXX - all assets for now - ALSO see maxa above excludes this


alternate 1 (PNS _) = PNS Asset
alternate 2 (PNS _) = PNS Liability

alternate 1 (PND _) = PND PGT
alternate 2 (PND _) = PND PLT

{-
-- To work, we need each alternate to make a tree of greater depth with at least one terminal node 
-- NullNum and NullNode are essentially production rules ...
-}

alternateSelection :: Gen a -> Mutator -> PNode -> Selector -> ST a PNode
alternateSelection g mut@(Mut nd _) (PNN (PNum x)) []
    = do num :: Int <- randRange g 0 99
         if num < tweakBias
          then do n' <- getDNum g nd -- i.e. tweak
                  return (PNN (PNum (x + n')))
          else alternateSelection1 g mut (PNN (PNum x)) []
-- replace PNNs with numbers 50% of the time
alternateSelection g mut@(Mut nd nnum) (PNN x) []
    = do num <- randRange g 0 1
         if num < 1
          then do n' <- getDNum g nnum
                  return (PNN (PNum n'))
          else alternateSelection1 g mut (PNN x) []
alternateSelection g mut (!n) [] 
    = alternateSelection1 g mut (n) [] 
alternateSelection g mut (!n) (!(num:r))
    = do new <- alternateSelection g mut (pick num n) r
         return (subst num new n)

alternateSelection1 g mut (!n) [] 
    = do num <- randRange g 1 max
         fixNullAndNum g mut (alternate num n) -- construct a tree, substing in all the Nulls with new trees
    where
       max = maxa n

fixNullAndNum :: Gen a -> Mutator -> PNode -> ST a (PNode)
-- fixNullAndNum _ _ p | (trace ("FIXING POLICY P : " ++ (show p)) False) = undefined
fixNullAndNum g (Mut nd newnumd) (PNN PNullVal) 
    = do n' <- getDNum g newnumd
         return (PNN (PNum n')) -- TODO maybe need a mutator parameter which is a bit smarter here ...
fixNullAndNum g mut nn@(PNN PNullNum) 
    = alternateSelection g mut nn [] 
fixNullAndNum g mut p@(PNP PNullPolicy)
    = alternateSelection g mut p [] 
fixNullAndNum g mut (!p)
    = do npicks <- mapM (fixSubTree g) picks
         return (newp npicks)
    where max = maxn p
          ixs = (enumFromTo 1 max)
          picks = map (\n -> pick n p) ixs
          fixSubTree ge pol = fixNullAndNum ge mut pol
          newp np = foldl (\pol (n,new) -> subst n new pol) p (zip ixs np)



type Selector = [Int]

selectNode :: Gen a -> PNode -> ST a Selector
-- selectNode _ p | (trace (" SN : " ++ (show p)) False) = undefined
selectNode g p
    = do num <- randRange g 0 99 -- 0 to 99 inclusive
         if num < haltMutationBias || maxn p == 0
          then return []
          else do rnum <- randRange g 1 max' -- 1 to max inclusive - a node selection
                  rest <- selectNode g (pick rnum p)
                  return (rnum:rest)
    where max' = max (maxn p) 1

substSelection :: PNode -> Selector -> PNode -> PNode
substSelection _ [] n = n
substSelection n (num:r) new 
    = let nnode = substSelection (pick num n) r new
      in subst num nnode n

getSelection :: PNode -> Selector -> PNode
getSelection n [] = n
getSelection n (num:r) = getSelection (pick num n) r


lookupDir :: PDirection -> (Double -> Double -> Bool)
lookupDir PLT = (<)
lookupDir PGT = (>)

mutate :: Gen a -> Mutator -> Management -> ST a Management 
mutate = mutatem 
mutatem :: Gen a -> Mutator -> Management -> ST a Management
mutatem g m (SimpleM p1 p2 d regs) =
        do p1' <- (mutatePolicy g m p1) 
           p2' <- (mutatePolicy g m p2)
           d' <- (mutatePNum g m d)
           regelem <- (mapM (mutatePolicy g m) regs)
           return (SimpleM p1' p2' d' regelem)


-- adds an if, with a randomly generated node in the condition
addIf :: Gen a -> Mutator -> Policy -> ST a Policy
addIf g m p = 
{-
do test <- (mutatePolicy g m PAccept)
                 p1' <- (mutatePolicy g m p) 
                 p2' <- (mutatePolicy g m p)
                 return (PIf test p1' p2') 
-}
-- simple version
    return (PIf p PAccept PReject)

-- stuffs up types 
{-
policyBreak :: Gen a -> Management -> ST a Management
policyBreak g (SimpleM p1 p2 d r) = toSM $ pd1 [ (PNP x) | x <- ([p1,p2]++r) ]
    where toSM (p1:(p2:r)) = (SimpleM (fromPNP p1) (fromPNP p2) d (map fromPNP r))
          pd1 ps = do ele <- randRange g 0 ((length ele)-1)
                      newele <- randRange 1 managementPolicyRegisters
                      let newNode = (PNN (PNum newele)
                      let oldtree = ps !! ele
                      sele <- selectNode g oldtree
                      let oldsubtree = getSelection oldtree sele
                          newtree = substSelection oldtree sele newNode
                      return $ substAt (newele+1) oldsubtree (substAt ele newtree ps)
          fromPNP (PNP x) = x
-}

mutateManyTest :: Gen a -> Mutator -> Management -> Int -> ST a [Management]
mutateManyTest g mut mgt i
               = do sequence (replicate i (mutate g mut mgt))

{-
Point mutation ...

we can do one of two things here -
- tweak the numbers in a policy
- change to an alternate policy

- XXX HARDCODED CHOICE BETWEEN TWEAKING AND FIDDLING THE TREE
-}

mutatePTree :: Gen a -> Mutator -> PNode -> ST a PNode 
mutatePTree g m p =
--    | (trace (" MP: " ++ (show p)) False ) = undefined
    do -- num <- randRange g 0 maxpol -- 0 to max *inclusive*
       asel <- selectNode g (p)
       apol <- alternateSelection g m (p) asel
       -- npol <- fixNullAndNumU m apol g 
       return apol

mutatePolicy :: Gen a -> Mutator -> Policy -> ST a Policy 
mutatePolicy g m p =
    do addIfTest <- randRange g 0 99
       if addIfTest < addIfBias
        then addIf g m p
        else mutatePTree g m (PNP p) >>= \x -> return $ fromPNP x
    where fromPNP (PNP p) = p

mutatePNum :: Gen a -> Mutator -> PNum -> ST a PNum 
mutatePNum g m p =
    mutatePTree g m (PNN p) >>= \x -> return $ fromPNN x
    where fromPNN (PNN p) = p